﻿// bc.cpp : 测试block cipher库
#include <cctype>
#include "../../common/encode/hex.h"
#include "../../common/include/std.h"
#include "../../common/net/socket.h"
#include "../../encrypt_and_encode/block_cipher/cb_cipher.h"
#include "../../encrypt_and_encode/block_cipher/cbc_mode.h"
#include "../../encrypt_and_encode/block_cipher/ctr_mode.h"
#include "../../encrypt_and_encode/block_cipher/fb_mode.h"
#include "../../encrypt_and_encode/block_cipher/gcm_mode.h"
#include "../../encrypt_and_encode/block_cipher/xor16.h"
// #ifdef LINUX
#include "openssl/evp.h"
// #endif

STD_BUFFER g_bufData;
BYTE g_bfKey[16];
BYTE g_bfIV[16];
BYTE g_bfEnc[256];
BYTE g_bfDec[256];

CCompositeBlockCipher g_clTest;

// #ifdef LINUX
#if 1
class C_SM4 : public IBlockCipherKernel
{
public:
    virtual int GetBlockLen() { return 16; };
    virtual int GetKeyLen() { return 16; };
    virtual int SetKey(BYTE *data, int iLen, void *pvResv = nullptr)
    {
        int iResult, iTmp;

        if(m_pstEnc == nullptr) m_pstEnc = EVP_CIPHER_CTX_new();
        if(m_pstDec == nullptr) m_pstDec = EVP_CIPHER_CTX_new();

        iResult = EVP_EncryptInit(m_pstEnc, EVP_sm4_ecb(), data, NULL);
        iResult = EVP_DecryptInit(m_pstDec, EVP_sm4_ecb(), data, NULL);

        EVP_CIPHER_CTX_set_padding(m_pstEnc, 0);
        EVP_CIPHER_CTX_set_padding(m_pstDec, 0);
        return 0;
    };

    virtual int Encrypt(BYTE *pIn, BYTE *pOut)
    {
        int iResult, iTmp = 16;

        iResult = EVP_EncryptUpdate(m_pstEnc, pOut, &iTmp, pIn, 16);

        if(iResult) return 0;
        else return -1;
    };

    virtual int Decrypt(BYTE *pIn, BYTE *pOut)
    {
        int iResult, iTmp = 16;

        iResult = EVP_DecryptUpdate(m_pstDec, pOut, &iTmp, pIn, 16);

        if(iResult) return 0;
        else return -1;
    };

protected:
    EVP_CIPHER_CTX *m_pstEnc = nullptr;
    EVP_CIPHER_CTX *m_pstDec = nullptr;
};
class C_AES : public IBlockCipherKernel
{
public:
    virtual int GetBlockLen() { return 16; };
    virtual int GetKeyLen() { return 16; };
    virtual int SetKey(BYTE *data, int iLen, void *pvResv = nullptr)
    {
        int iResult, iTmp;

        if(m_pstEnc == nullptr) m_pstEnc = EVP_CIPHER_CTX_new();
        if(m_pstDec == nullptr) m_pstDec = EVP_CIPHER_CTX_new();

        iResult = EVP_EncryptInit(m_pstEnc, EVP_aes_128_ecb(), data, NULL);
        iResult = EVP_DecryptInit(m_pstDec, EVP_aes_128_ecb(), data, NULL);

        EVP_CIPHER_CTX_set_padding(m_pstEnc, 0);
        EVP_CIPHER_CTX_set_padding(m_pstDec, 0);
        return 0;
    };

    virtual int Encrypt(BYTE *pIn, BYTE *pOut)
    {
        int iResult, iTmp = 16;

        iResult = EVP_EncryptUpdate(m_pstEnc, pOut, &iTmp, pIn, 16);

        if(iResult) return 0;
        else return -1;
    };

    virtual int Decrypt(BYTE *pIn, BYTE *pOut)
    {
        int iResult, iTmp = 16;

        iResult = EVP_DecryptUpdate(m_pstDec, pOut, &iTmp, pIn, 16);

        if(iResult) return 0;
        else return -1;
    };

protected:
    EVP_CIPHER_CTX *m_pstEnc = nullptr;
    EVP_CIPHER_CTX *m_pstDec = nullptr;
};
#endif

int TestGCM()
{
    int iResult, iTmp;
    C_GCM_MODE clGCM;
    int iWrite = 0;
    string strOut;

    g_clTest.m_pobjKernel->SetKey(g_bfKey, 16);

    iResult = clGCM.Init(g_bfIV, 12, nullptr, 0, g_clTest.m_pobjKernel);
    cout << "gcm.init ret " << iResult << endl;

    iTmp    = 256;
    iResult = clGCM.Encrypt(g_bufData.data(), g_bufData.size(), g_bfEnc, &iTmp);
    // cout << "gcm.enc ret " << iResult << endl;
    iWrite += iTmp;

    iTmp    = 256 - iWrite;
    iResult = clGCM.Encrypt(nullptr, 0, g_bfEnc + iWrite, &iTmp);
    // cout << "gcm.enc ret " << iResult << endl;
    iWrite += iTmp;

    HexEncode(g_bfEnc, iWrite, strOut);
    cout << "gcm.enc(" << iWrite << ") = " << strOut << endl;

    iWrite  = 0;
    iResult = clGCM.Init(g_bfIV, 12, nullptr, 0, g_clTest.m_pobjKernel);
    cout << "gcm.init ret " << iResult << endl;

    iTmp    = 256;
    iResult = clGCM.Decrypt(g_bfEnc, g_bufData.size(), g_bfDec, &iTmp);
    // cout << "gcm.dec ret " << iResult << endl;
    iWrite += iTmp;

    iTmp    = 256 - iWrite;
    iResult = clGCM.Decrypt(nullptr, 0, g_bfDec + iWrite, &iTmp);
    // cout << "gcm.dec ret " << iResult << endl;
    iWrite += iTmp;

    HexEncode(g_bfDec, iWrite, strOut);
    cout << "gcm.dec(" << iWrite << ") = " << strOut << endl;

    return 0;
}

int Test(const char *pszName, C_ECB_MODE *pobjMode)
{
    int iResult, iTmp;
    BYTE bfIV[16];
    string strOut;

    cout << pszName << endl;
    g_clTest.m_pobjMode = pobjMode;

    memcpy(bfIV, g_bfIV, 16);
    iTmp    = 256;
    iResult = g_clTest.Encrypt(g_bfKey, 16, bfIV, 16, g_bufData.data(), g_bufData.size(), g_bfEnc, &iTmp);
    HexEncode(g_bfEnc, g_bufData.size(), strOut);
    cout << "enc = " << strOut << endl;

    memcpy(bfIV, g_bfIV, 16);
    iTmp    = 256;
    iResult = g_clTest.Decrypt(g_bfKey, 16, bfIV, 16, g_bfEnc, g_bufData.size(), g_bfDec, &iTmp);
    HexEncode(g_bfDec, g_bufData.size(), strOut);
    cout << "dec = " << strOut << endl;

    DELETE_AND_NULL(g_clTest.m_pobjMode);
    return 0;
}

int main()
{
    int iResult, iTmp;
    string strData, strOut;
    int iDataLen = 32;

    // for (int i = 0; i < 16; i++) g_bfKey[i] = i + 1;
    // HexDecode("2b7e151628aed2a6abf7158809cf4f3c", -1, g_bfKey);
    bzero(g_bfKey, 16);

    // for (int i = 0; i < 256; i++) g_bfData[i] = i;
    // strData = "6BC1BEE22E409F96E93D7E117393172AAE2D8A571E03AC9C9EB76FAC45AF8E51";
    strData = "00000000000000000000000000000000";
    HexDecode(strData, g_bufData);
    cout << "data = " << strData << endl;

    bzero(g_bfIV, 16);
    // for(int i = 0; i < 16; i++) g_bfIV[i] = i;
// #ifdef LINUX
#if 1
    // g_clTest.m_pobjKernel = new C_AES();
    g_clTest.m_pobjKernel = new C_SM4();
#else
    g_clTest.m_pobjKernel = new C_XOR16();
#endif
    TestGCM();
    
    iResult = Test("ecb test", new C_ECB_MODE());
    iResult = Test("cbc test", new C_CBC_MODE());
    iResult = Test("cfb test", new C_CFB_MODE());
    iResult = Test("ofb test", new C_OFB_MODE());

    HexDecode("f0f1f2f3f4f5f6f7f8f9fafbfcfdfeff", -1, g_bfIV);
    iResult = Test("ctr test", new C_CTR_MODE());
    
    cout << "end" << endl;
    return 0;
}
